#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
极简加密聊天 v3.0 - 稳定版本
修复连接问题，适用于所有场景（同一台电脑、不同设备、跨网络）
"""

import socket
import threading
import time
import secrets
import sys
import os
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.kdf.hkdf import HKDF


class SimpleCrypto:
    """极简加密类 - 使用P-384曲线提供192位安全级别"""
    
    def __init__(self):
        print("[密钥] 生成P-384加密密钥...")
        # 生成ECDH密钥对 - P-384
        self.private_key = ec.generate_private_key(ec.SECP384R1())
        self.public_key = self.private_key.public_key()
        self.shared_key = None
        self.aes = None
    
    def get_public_key_bytes(self):
        """获取公钥字节 - P-384需要48字节 x + 48字节 y"""
        return self.public_key.public_numbers().x.to_bytes(48, 'big') + \
               self.public_key.public_numbers().y.to_bytes(48, 'big')
    
    def compute_shared_key(self, peer_public_key_bytes):
        """计算共享密钥 - 处理P-384的96字节公钥"""
        try:
            # 重构对方公钥 (48字节 x + 48字节 y)
            x = int.from_bytes(peer_public_key_bytes[:48], 'big')
            y = int.from_bytes(peer_public_key_bytes[48:], 'big')
            peer_public_numbers = ec.EllipticCurvePublicNumbers(x, y, ec.SECP384R1())
            peer_public_key = peer_public_numbers.public_key()
            
            # 计算共享密钥
            shared_secret = self.private_key.exchange(ec.ECDH(), peer_public_key)
            
            # 派生AES密钥
            derived_key = HKDF(
                algorithm=hashes.SHA256(),
                length=32,
                salt=None,
                info=b'terminal_chat_v3_p384'
            ).derive(shared_secret)
            
            self.aes = AESGCM(derived_key)
            print("[成功] 密钥交换成功！使用P-384+AES-256-GCM加密")
            return True
        except Exception as e:
            print(f"[错误] 密钥交换失败: {e}")
            return False
    
    def encrypt_message(self, message):
        """加密消息"""
        if not self.aes:
            return None
        nonce = secrets.token_bytes(12)
        ciphertext = self.aes.encrypt(nonce, message.encode('utf-8'), None)
        return nonce + ciphertext
    
    def decrypt_message(self, encrypted_data):
        """解密消息"""
        if not self.aes or len(encrypted_data) < 12:
            return None
        nonce = encrypted_data[:12]
        ciphertext = encrypted_data[12:]
        try:
            plaintext = self.aes.decrypt(nonce, ciphertext, None)
            return plaintext.decode('utf-8')
        except:
            return None


class TerminalChat:
    """终端聊天应用"""
    
    def __init__(self):
        self.crypto = SimpleCrypto()
        self.socket = None
        self.connected = False
        self.server_socket = None
        self.running = True
        self.receive_thread = None
        self.server_thread = None
        
        # 打印启动信息
        self.print_header()
        
    def print_header(self):
        """打印程序头部信息"""
        print("=" * 60)
        print("  极简加密聊天 v3.0 - 稳定版本")
        print("  P-384椭圆曲线 + AES-256-GCM加密")
        print("  IPv6点对点直连通信 (支持所有场景)")
        print("=" * 60)
        
    def get_my_ipv6(self):
        """获取本机IPv6地址"""
        try:
            s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
            s.connect(("2001:4860:4860::8888", 80))  # Google DNS IPv6
            my_ipv6 = s.getsockname()[0]
            s.close()
            return my_ipv6
        except:
            return "::1"  # 回退到IPv6回环地址
    
    def start_server(self):
        """启动服务器监听"""
        def server_thread():
            try:
                self.server_socket = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
                self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                # 绑定到所有IPv6接口
                self.server_socket.bind(('::', 8888))
                self.server_socket.listen(1)
                server_addr = self.server_socket.getsockname()
                print(f"[监听] 在端口 {server_addr[1]} 监听所有IPv6接口，等待连接...")
                
                while self.running and not self.connected:
                    try:
                        # 设置accept超时，避免无限阻塞
                        self.server_socket.settimeout(1.0)
                        client_socket, addr = self.server_socket.accept()
                        
                        if not self.connected:
                            print(f"[连接] 收到来自 {addr[0]}:{addr[1]} 的连接")
                            self.socket = client_socket
                            # 不要在这里退出循环，让连接处理在单独线程中进行
                            threading.Thread(target=self.handle_connection, args=(True,), daemon=True).start()
                            return  # 连接处理开始后退出服务器线程
                        else:
                            client_socket.close()
                            
                    except socket.timeout:
                        continue  # 超时是正常的，继续监听
                    except Exception as e:
                        if self.running:
                            print(f"[错误] 接受连接失败: {e}")
                        break
                        
            except Exception as e:
                print(f"[错误] 服务器启动失败: {e}")
        
        self.server_thread = threading.Thread(target=server_thread, daemon=True)
        self.server_thread.start()
    
    def connect_to_peer(self, address):
        """连接到对方"""
        try:
            print(f"[连接] 正在连接到 {address}...")
            
            # 如果是回环地址，稍微延迟一下确保服务器就绪
            if address in ["::1", "127.0.0.1"] or address.startswith("240e:"):
                time.sleep(0.5)
            
            self.socket = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
            self.socket.connect((address, 8888))
            # 显示本地端口信息
            local_addr = self.socket.getsockname()
            print(f"[本地] 使用端口 {local_addr[1]} 连接到 {address}:8888")
            threading.Thread(target=self.handle_connection, args=(False,), daemon=True).start()
            return True
        except Exception as e:
            print(f"[错误] 连接失败: {e}")
            return False
    
    def handle_connection(self, is_server):
        """处理连接和密钥交换"""
        try:
            print("[握手] 开始密钥交换...")
            
            # 密钥交换
            my_public_key = self.crypto.get_public_key_bytes()
            
            # 设置socket超时
            self.socket.settimeout(5.0)
            
            if is_server:
                # 服务器：先接收，后发送
                peer_public_key = self.socket.recv(96)  # P-384需要96字节
                self.socket.send(my_public_key)
            else:
                # 客户端：先发送，后接收
                self.socket.send(my_public_key)
                peer_public_key = self.socket.recv(96)  # P-384需要96字节
            
            # 计算共享密钥
            if self.crypto.compute_shared_key(peer_public_key):
                self.connected = True
                # 显示连接详细信息
                local_addr = self.socket.getsockname()
                peer_addr = self.socket.getpeername()
                print(f"[连接] 本地: {local_addr[0]}:{local_addr[1]} <-> 对方: {peer_addr[0]}:{peer_addr[1]}")
                print("[成功] 连接成功！现在可以安全聊天了")
                print("[提示] 输入消息按回车发送，输入 '/quit' 退出")
                print("-" * 60)
                
                # 清除socket超时
                self.socket.settimeout(None)
                
                # 启动接收消息线程
                self.start_receive_thread()
                
                # 不在连接线程中调用input_loop，让主线程处理输入
            else:
                print("[错误] 密钥交换失败")
                
        except Exception as e:
            print(f"[错误] 连接处理失败: {e}")
            self.connected = False
    
    def start_receive_thread(self):
        """启动接收消息线程"""
        def receive_thread():
            while self.connected and self.running:
                try:
                    # 阻塞等待消息，不设置超时
                    self.socket.settimeout(None)
                    # 接收消息长度
                    length_data = self.socket.recv(4)
                    if not length_data:
                        print("[信息] 对方关闭了连接")
                        break
                    
                    msg_length = int.from_bytes(length_data, 'big')
                    if msg_length <= 0 or msg_length > 1000000:  # 防止异常长度
                        print(f"[警告] 接收到异常消息长度: {msg_length}")
                        break
                    
                    # 接收加密消息
                    encrypted_data = b''
                    while len(encrypted_data) < msg_length:
                        chunk = self.socket.recv(msg_length - len(encrypted_data))
                        if not chunk:
                            print("[信息] 消息接收不完整")
                            break
                        encrypted_data += chunk
                    
                    if len(encrypted_data) != msg_length:
                        print(f"[警告] 消息长度不匹配: 期望{msg_length}, 实际{len(encrypted_data)}")
                        continue
                    
                    # 解密消息
                    message = self.crypto.decrypt_message(encrypted_data)
                    if message:
                        timestamp = time.strftime("%H:%M:%S")
                        print(f"\r[{timestamp}] 对方: {message}")
                        print("我: ", end="", flush=True)
                    else:
                        print("[警告] 消息解密失败")
                    
                except Exception as e:
                    print(f"[错误] 接收消息出错: {e}")
                    break
            
            # 连接断开
            if self.connected:
                self.connected = False
                print(f"\n[断开] 连接已断开")
        
        self.receive_thread = threading.Thread(target=receive_thread, daemon=True)
        self.receive_thread.start()
    
    def send_message(self, message):
        """发送消息"""
        if not self.connected:
            return False
        
        try:
            # 加密消息
            encrypted_data = self.crypto.encrypt_message(message)
            if encrypted_data:
                # 发送消息长度 + 加密消息
                self.socket.send(len(encrypted_data).to_bytes(4, 'big'))
                self.socket.send(encrypted_data)
                return True
            else:
                print("[错误] 消息加密失败")
                return False
        except Exception as e:
            print(f"[错误] 发送失败: {e}")
            return False
    
    def input_loop(self):
        """输入循环"""
        while self.connected and self.running:
            try:
                message = input("我: ").strip()
                
                if message == '/quit':
                    print("[退出] 准备退出...")
                    break
                elif message:
                    if self.send_message(message):
                        # 发送成功，显示带时间戳的消息
                        timestamp = time.strftime("%H:%M:%S")
                        print(f"\r[{timestamp}] 我: {message}")
                        print("我: ", end="", flush=True)
                    else:
                        print("[错误] 发送失败，连接可能已断开")
                        break
                        
            except (KeyboardInterrupt, EOFError):
                print("\n[中断] 用户中断")
                break
            except Exception as e:
                print(f"[错误] 输入循环出错: {e}")
                break
        
        self.cleanup()
    
    def cleanup(self):
        """清理资源"""
        print("[清理] 正在清理资源...")
        self.running = False
        self.connected = False
        
        if self.socket:
            try:
                self.socket.close()
            except:
                pass
            
        if self.server_socket:
            try:
                self.server_socket.close()
            except:
                pass
                
        print("[完成] 程序已退出，聊天记录已清除")
    
    def run(self):
        """运行聊天程序"""
        try:
            # 显示本机IPv6地址
            my_ipv6 = self.get_my_ipv6()
            print(f"[地址] 我的IPv6地址: {my_ipv6}")
            print("[提示] 对方需要使用此地址连接到您")
            
            # 启动服务器监听
            self.start_server()
            
            # 等待服务器启动
            time.sleep(0.5)
            
            # 用户选择模式
            print("\n请选择:")
            print("1. 等待对方连接我")
            print("2. 连接对方")
            
            while True:
                choice = input("请输入选择 (1/2): ").strip()
                
                if choice == '1':
                    print("[等待] 等待对方连接...")
                    # 服务器模式，等待连接并处理
                    while not self.connected and self.running:
                        time.sleep(0.1)
                    # 连接建立后，在主线程中处理输入
                    if self.connected:
                        self.input_loop()
                    break
                    
                elif choice == '2':
                    peer_address = input("请输入对方IPv6地址: ").strip()
                    if peer_address:
                        if self.connect_to_peer(peer_address):
                            # 等待连接完成
                            timeout = 10
                            while not self.connected and self.running and timeout > 0:
                                time.sleep(0.1)
                                timeout -= 1
                            if self.connected:
                                # 连接成功，在主线程中处理输入
                                self.input_loop()
                                break
                            else:
                                print("[错误] 连接超时，请重试")
                        else:
                            print("[错误] 连接失败，请重试")
                    
                else:
                    print("请输入1或2")
                    
        except KeyboardInterrupt:
            print("\n\n[退出] 程序被中断")
            self.cleanup()


def main():
    """主函数"""
    try:
        chat = TerminalChat()
        chat.run()
    except Exception as e:
        print(f"程序错误: {e}")


if __name__ == "__main__":
    main() 